import time
from datetime import datetime
from io import StringIO
from typing import List, Tuple

import taosws
from sqlalchemy import CursorResult


class TDengineClientWithWS:
    def __init__(self, username, password, host, port, db):
        self.username = username
        self.password = password
        self.host = host
        self.port = port
        self.db = db
        self.conn = taosws.connect(f"taosws://{self.username}:{self.password}@{self.host}:{self.port}")

    def disconnect(self):
        if self.conn:
            self.conn.close()

    def _execute(self, sql):
        while True:
            try:
                return self.conn.execute(sql)
            except Exception as e:
                print(e)
                time.sleep(3)
                continue

    @staticmethod
    def rows_to_dict(result: CursorResult, one=False):
        """
        将结果集转成字典
        """
        columns = tuple(map(str, result.keys()))
        rows = result.fetchone() if one else result.fetchall()
        if one:
            return dict(zip(columns, rows))
        return [dict(zip(columns, row)) for row in rows]

    def create_super_table(self,
                           stable: str,
                           ts_field: str,
                           metric_fields: List,
                           tags: List):
        """
        创建超级表
        @stable: 超级表的名称
        @ts_field: 时间戳字段的名称
        @metric_fields: 采集的物理量对应的字段名和类型
        @tags: 标签字段和类型
        """
        if type(metric_fields) is not list:
            raise ValueError("metric_fields 必须是一个列表，里面包含二元组")
        if type(tags) is not list:
            raise ValueError("tags 必须是一个列表，里面包含二元组")
        ddl = f"""
                CREATE STABLE {stable} (
                    {ts_field} timestamp,
                    {', '.join([' '.join([metric, tp]) for metric, tp in metric_fields])}
                ) TAGS (
                    {', '.join([' '.join([tag, tp]) for tag, tp in tags])}
                )
            """
        try:
            self._execute(ddl)
        except Exception as e:
            print(f"超级表建立失败, 错误原因: {e}")
            raise e

    def create_child_table(self, table: str, stable: str, tags_value: Tuple):
        """
        创建子表，必须指定一张超级表
        @table: 子表的名称
        @stable: 超级表的名称
        @tags_value: 标签值, 一个二元组
        """
        ddl = f"""
                CREATE TABLE {table} USING {stable}
                TAGS {tags_value}
            """
        try:
            self._execute(ddl)
        except Exception as e:
            print(f"子表建立失败, 错误原因: {e}")
            raise e

    def td_insert_data(self,
                       table: str,
                       values: List[Tuple],
                       columns: List = None):
        """
        往 tdengine 批量插入数据
        @table: 表名
        @value: 数据
        @columns: 列名
        """
        buf = StringIO()
        for value in values:
            value = list(
                str(item) if isinstance(item, datetime) else "null" if item is None else item
                for item in value
            )
            buf.write(f"({str(value)[1: -1]})")

        query = f"""
                INSERT INTO {self.db}.{table} {'(' + ','.join(columns) + ')' if columns is not None else ""} 
                VALUES {buf.getvalue()}
            """
        buf.truncate(0)
        buf.close()
        try:
            self._execute(query)
        except Exception as e:
            print(f"数据插入失败, 错误原因: {e}")
            raise e


td_client = TDengineClientWithWS(host='192.168.10.130', port=6041, username='root', password='taosdata', db='monitor')
# td_client.connect()
td_client._execute(f'use {td_client.db}')
# td_client.create_super_table('pidmonitor', 'ts',
#                              [('name', 'NCHAR(16)'), ('status', 'NCHAR(16)'), ('cpu_percent', 'float'),
#                               ('memory_percent', 'float'), ('rss', 'float')], [('groupId', 'INT')])


